#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File     : loader.py
@Project  : pipecoco
@Date     : 2021/8/13
@Author   : Zhang Jinyang
@Contact  : zhang-jy@sjtu.edu.cn
'''

from comm import pipecoco_pb2
import numpy as np
import lz4.frame

def pack_data(model_name, x, i):
    """
    将数据封装成grpc类，将数据以二进制形式发送
    """
    fragment_shape = list(x.shape)

    if len(x.tobytes())==0:
        compressed = x.tobytes()
    else:
        compressed = x.tobytes()
        # compressed = lz4.frame.compress(x.tobytes())
    fragment_info = pipecoco_pb2.block_data(model_name=model_name, block_id=i, fragment=compressed,
                                            fragment_shape=fragment_shape)
    return fragment_info


def unpack_data(request):
    """
    对收到数据进行解包
    """

    fragment_shape = request.fragment_shape
    if len(request.fragment)==0:
        x = request.fragment
    else:
        x = request.fragment
        # x = lz4.frame.decompress(request.fragment)
    x = np.frombuffer(x, dtype=np.float32)
    x = x.reshape(fragment_shape)
    return x
